#ifndef __MMS_WATCH_POINT__
#define __MMS_WATCH_POINT__

#include <map>

namespace mm
{
	struct iwatcher_t : RefalbeImp
	{
		virtual size_t	before_call( size_t uid, hk::exec_context * pctx, size_t deny_count = false ) { return true; };
		virtual void	after_call( size_t uid, hk::exec_context * pctx ) {};
	};

	struct item_t
	{
		size_t				uid;
		refp<iwatcher_t>	watcher;
	};

	bool need_handle_this_hook( hk::exec_context * pctx, bool _after );

	struct watch_point_t : hk::hookfn_t, RefalbeImp
	{
		typedef hk::hookfn_t			super_t;
		typedef std::vector< item_t >	watcher_map_t;

		watcher_map_t	_watchers;
		mutex_t			_mutex;

		//////////////////////////////////////////////////////////////////////////

		watch_point_t()
		{
			super_t::user_data = this;
		}
		~watch_point_t()
		{
			sync_t __sync__(_mutex);
			_watchers.clear();
		}

		static void * operator new ( size_t c )
		{
			return hk::codeheap_alloc( c );
		}

		static void operator delete ( void * p )
		{
			return hk::codeheap_free( p );
		}

		long init( const void * fnp, size_t retn, bool patch_now = false )
		{
			return __super::init( fnp, retn, before_call__, after_call__, patch_now );
		}

		bool actived()
		{
			return __super::patched != 0;
		}

		long attach( iwatcher_t * p, size_t ud )
		{
			if( !p ) return E_INVALIDARG;
			sync_t __sync__(_mutex);
			try 
			{
				item_t item;
				item.uid = ud;
				item.watcher = p;
				_watchers.push_back(item);
				return S_OK;
			}
			catch(...)
			{
				
			}
			return E_FAIL;
		}

		long detach( iwatcher_t * p )
		{
			if( !p ) return E_INVALIDARG;
			sync_t __sync__(_mutex);
			try 
			{
				watcher_map_t::iterator it = _watchers.begin();
				for( ; it != _watchers.end(); )
				{
					if( it->watcher == p )
						it = _watchers.erase( it );
					else
						++ it;
				}
			}
			catch( ... )
			{
				return E_FAIL;	
			}
			return S_OK;
		}

		const void * hook_point()
		{
			return super_t::hooked;
		}
		const void * hook_target()
		{
			return super_t::get_hook_target();
		}

		void destroy()
		{
			sync_t __sync__(_mutex);
			_watchers.clear();
			super_t::disable();
		}

	protected:

		size_t before_call( hk::exec_context * pctx )
		{
			sync_t __sync__(_mutex);
			watcher_map_t::iterator it = _watchers.begin();
			size_t deny = 0;
			for( ; it != _watchers.end(); ++ it )
			{
				item_t & item = *it;
				try 
				{
					size_t f_pass = item.watcher->before_call( item.uid, pctx, deny );
					if( !f_pass ) ++ deny;
				}
				catch(...)
				{
				}
			}
			return deny ? false : true;
		}

		void after_call( hk::exec_context * pctx )
		{
			sync_t __sync__(_mutex);
			watcher_map_t::iterator it = _watchers.begin();
			for( ; it != _watchers.end(); ++ it )
			{
				item_t & item = *it;
				try 
				{
					item.watcher->after_call( item.uid, pctx );
				}
				catch(...)
				{
				}
			}
		}
		
		static size_t before_call__( hk::exec_context * pctx )
		{
			try 
			{
				if( !need_handle_this_hook(pctx, false) ) return true;
				hk::hookfn_t* hkp = (hk::hookfn_t*)pctx->self;
				return ((watch_point_t*)(hkp->user_data))->before_call(pctx);
			}
			catch( ... )
			{

			}
			return true;
		};

		static void after_call__( hk::exec_context * pctx )
		{
			try 
			{
				if( !need_handle_this_hook(pctx, true) ) return ;
				hk::hookfn_t* hkp = (hk::hookfn_t*)pctx->self;
				((watch_point_t*)(hkp->user_data))->after_call( pctx );
			}
			catch( ... )
			{

			}
		};
	};

	typedef std::set<size_t>	tid_map_t;
	typedef std::set<size_t>	mba_map_t;

	enum
	{
		FILTER_BY_NONE		= 0,
		FILTER_BY_TID		= 1,
		FILTER_BY_MEM		= 2,
	};

	struct wpfilter_t
	{
		size_t		_mon_mask;
		tid_map_t	_mon_threads;
		mba_map_t	_mon_mblocks;
		mutex_t		_lock;

		wpfilter_t()
			: _mon_mask(FILTER_BY_NONE)
		{

		}

		size_t set_filter_mask( size_t add, size_t del )
		{
			sync_t __lock__(_lock);
			_mon_mask &= ~del;
			_mon_mask |= add;
			return _mon_mask;
		}

		long watch_thread( size_t tid, bool f_mon )
		{
			sync_t __lock__(_lock);
			if( f_mon ) _mon_threads.insert( tid );
			else _mon_threads.erase( tid );
			return S_OK;
		}

		long watch_memory( const void * pba, bool f_mon )
		{
			sync_t __lock__(_lock);
			MEMORY_BASIC_INFORMATION mbi;
			if( !VirtualQuery( pba, &mbi, sizeof(mbi) ) ) return E_FAIL;
			_mon_mblocks.insert( (size_t)mbi.AllocationBase );
			return S_OK;
		}

		bool need_handle( hk::exec_context * pctx, bool f_after_call = false )
		{
			sync_t __lock__(_lock);

			if( _mon_mask == FILTER_BY_NONE )
				return true;

			if( _mon_mask & FILTER_BY_TID )
			{
				tid_map_t::iterator it = _mon_threads.find( GetCurrentThreadId() );
				if( it == _mon_threads.end() )
					return false;
			}
			if( _mon_mask & FILTER_BY_MEM )
			{
				void * frames[16] = {};
				size_t cf = hk::walk_frame( pctx, frames, 16, f_after_call );
				if( cf > 16) __asm int 3;
				size_t i = 0;
				for( ; i < cf; ++ i )
				{
					MEMORY_BASIC_INFORMATION mbi;
					if( !VirtualQuery( frames[i], &mbi, sizeof(mbi) ) )
						continue;

					mba_map_t::iterator it = _mon_mblocks.find( (size_t)mbi.AllocationBase );
					if( it != _mon_mblocks.end() )
						break;
				}
				if( i == cf )
					return false;
			}
			return true;
		}
	};

	//////////////////////////////////////////////////////////////////////////

	class wp_mgmt_t : public wpfilter_t
	{
	protected:

		typedef std::map< const void *, refp<watch_point_t> >	wp_map_t;

		wp_map_t				_wps;
		mutex_t					_mutex;

	public:

		wp_mgmt_t()
		{
		}
		~wp_mgmt_t()
		{
			sync_t __sync__(_mutex);
			_wps.clear();
		}

		long set_watch_point( const void * fep, size_t retn, bool enable_now = false )
		{
			sync_t __sync__(_mutex);
			refp<watch_point_t> wp;
			long lr = _set_watch_point( fep, retn, wp.pp() );
			if( lr < 0 ) return lr;
			if( enable_now )
				wp_enable( wp, enable_now );
			return S_OK;
		}

		long wp_enable( watch_point_t * wp, bool f )
		{
			sync_t __sync__(_mutex);
			long lr = E_FAIL;
			// suspand other threads
			if( f ) 
				lr = wp->enable();
			else
				lr = wp->disable();
			// resume other threads
			return lr;
		}

		long banch_active( bool active )
		{
			sync_t __sync__(_mutex);
			// suspand other threads
			wp_map_t::iterator it = _wps.begin();
			for( ; it != _wps.end(); ++ it )
			{
				if( active )
					it->second->enable();
				else
					it->second->disable();
			}
			// resume other threads
			return S_OK;
		}

		long attach( const void * p, iwatcher_t * wo, void* ud = 0, bool force = false, size_t retn = 0 )
		{
			return attach( p, wo, (size_t)ud, force, retn );
		}

		long attach( const void * p, iwatcher_t * wo, size_t ud = 0, bool force = false, size_t retn = 0 )
		{
			sync_t __sync__(_mutex);
			refp<watch_point_t> wp;
			wp_map_t::iterator it = _wps.find( p );
			if( it == _wps.end() )
			{
				if( !force ) return E_FAIL;
				long lr = _set_watch_point( p, retn, wp.pp() );
				if( lr < 0 ) return lr;
			}
			else
			{
				wp = it->second;
			}
			if( !wp ) return E_FAIL;
			return wp->attach( wo, ud );
		}

		long destroy( const void * p )
		{
			sync_t __sync__(_mutex);
			wp_map_t::iterator it = _wps.find( p );
			if( it == _wps.end() ) return E_FAIL;
			refp<watch_point_t> wp = it->second;
			_wps.erase( it );
			wp->destroy();
			return 0;
		}

		long detach( const void * p, iwatcher_t * wo )
		{
			sync_t __sync__(_mutex);
			refp<watch_point_t> wp;
			wp_map_t::iterator it = _wps.find( p );
			if( it == _wps.end() ) return E_FAIL;
			wp = it->second;
			if( !wp ) return E_FAIL;
			return wp->detach( wo );
		}

		long active_buildin()
		{
			sync_t __sync__(_mutex);

			//buildin_events = new buildin_watcher_t();
			//if( !buildin_events ) return E_OUTOFMEMORY;

			//if( 0 )// _ki_user_exception_dispatcher )
			//{
			//	_ki_user_exception_dispatcher->active();
			//	_ki_user_exception_dispatcher->connect( buildin_events, buildin_watcher_t::api_ki_user_exception_dispatcher );
			//}

			//HMODULE mod = GetModuleHandleW(L"ntdll.dll");
			//if( !mod ) return E_FAIL;

			//void * lpNtAllocateVirtualMemory = GetProcAddress( mod, "NtAllocateVirtualMemory" );
			//if( lpNtAllocateVirtualMemory ) 
			//	attach( lpNtAllocateVirtualMemory, buildin_events, buildin_watcher_t::api_nt_allocate_virtual_memory, true, 6 * 4 );

			//void * lpNtProtectVirtualMemory = GetProcAddress( mod, "NtProtectVirtualMemory" );
			//if( lpNtProtectVirtualMemory ) 
			//	attach( lpNtProtectVirtualMemory, buildin_events, buildin_watcher_t::api_nt_protect_virtual_memory, true, 5 * 4 );

			return S_OK;
		}

	protected:

		long _set_watch_point( const void * fep, size_t retn, watch_point_t ** ppwp )
		{
			refp<watch_point_t> wp;
			wp_map_t::iterator it = _wps.find( fep );
			if( it != _wps.end() )
			{
				wp = it->second;
			}
			else
			{
				wp = new watch_point_t();
				if( !wp ) return E_OUTOFMEMORY;

				long lr = wp->init( fep, retn );
				if( lr < 0 ) return lr;
				
				_wps[fep] = wp;
			}

			*ppwp = wp.detach();

			return S_OK;
		}
	};

	static mm::wp_mgmt_t	ProcessWatchPointManamger;

	long wpm_trace_thread( size_t tid )
	{
		return 0;
	}

	long wpm_trace_memory( void * p )
	{
		return 0;	
	}

	long wpm_set_watch_point( const void * fep, size_t retn, bool enable_now = false )
	{
		return ProcessWatchPointManamger.set_watch_point( fep, retn, enable_now );
	}

	long wpm_banch_active( bool f_active )
	{
		return ProcessWatchPointManamger.banch_active( f_active );
	}

	long wpm_attach( const void * p, iwatcher_t * wo, void* ud = 0, bool force = false, size_t retn = 0 )
	{
		return ProcessWatchPointManamger.attach( p, wo, ud, force, retn );
	}

	long wpm_attach( const void * p, iwatcher_t * wo, size_t ud = 0, bool force = false, size_t retn = 0 )
	{
		return ProcessWatchPointManamger.attach( p, wo, ud, force, retn );
	}

	long wpm_detach( const void * p, iwatcher_t * wo )
	{
		return ProcessWatchPointManamger.detach( p, wo );
	}
	long wpm_destroy( const void * p)
	{
		return ProcessWatchPointManamger.destroy( p );
	}
	bool need_handle_this_hook( hk::exec_context * pctx, bool _after )
	{
		return ProcessWatchPointManamger.need_handle( pctx, _after );
	}
};

#endif
